package com.wolf.antlr;

import com.wolf.antlr.python.PythonLexer;
import com.wolf.antlr.python.PythonParser;
import com.wolf.antlr.python.PythonParserBaseListener;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.TokenStream;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.ParseTreeWalker;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class PythonSearcher {
    static List<MethodFragment> searchFunctions(String fileName, PythonListenerImp listener) {
        Lexer lexer;
        try {
            lexer = new PythonLexer(CharStreams.fromFileName(fileName));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        TokenStream tokenStream = new CommonTokenStream(lexer);
        PythonParser parser = new PythonParser(tokenStream);
        ParseTree tree = parser.root();
        ParseTreeWalker walker = new ParseTreeWalker();
        walker.walk(listener, tree);
        return listener.getMethodFragments();
    }

    public static void main(String[] args) {
        PythonListenerImp listener = new PythonListenerImp();
        String fileName = Main.parserArgs(args, listener);
        List<MethodFragment> methodFragments = searchFunctions(fileName, listener);
        methodFragments.forEach(methodFragment -> System.out.printf("%s, %s, %s, %s, %s\n",
                methodFragment.startIndex(),
                methodFragment.endIndex(),
                methodFragment.startLine(),
                methodFragment.endLine(),
                methodFragment.methodDeclaration()));
    }

    static class PythonListenerImp extends PythonParserBaseListener implements IListener {
        private String getMethodDeclaration(PythonParser.FuncdefContext ctx) {
            // 获取函数名
            String func_name = ctx.name().getText();
            // 获取参数列表
            List<String> params = new ArrayList<>();
            if (ctx.typedargslist() != null) {
                for (PythonParser.Def_parametersContext param : ctx.typedargslist().def_parameters()) {
                    params.add(param.getText().replace(",", ", "));
                }
            }
            // 拼接函数声明的完整格式
            StringBuilder func_decl = new StringBuilder();
            func_decl.append("def ");
            func_decl.append(func_name);
            func_decl.append("(");
            for (String param : params) {
                func_decl.append(param);
            }
            func_decl.append(")");
            func_decl.append(":");
            return func_decl.toString();
        }

        @Override
        public void enterFuncdef(PythonParser.FuncdefContext ctx) {
            super.enterFuncdef(ctx);
            String methodDeclaration = this.getMethodDeclaration(ctx);
            this.search(ctx, methodDeclaration);
        }
    }
}
